#include <iostream>
#include "postprocess.h"
#include "preprocess.h"
#include "yolov8.h"
#include "videofile.h"
#include <yolo_threadpool.h>
#include "time_duration.h"
#include "run_once.h"

#include <opencv2/opencv.hpp>

int main(int argc, char const *argv[])
{
    std::string model_path = "/home/orangepi/Code/inference-rknn/resources/yolov8/yolov8n.rknn";
    std::string input_path = "/home/orangepi/Code/rknn-bytetrack/videos/palace.mp4";
    std::string labels_path = "/home/orangepi/Code/inference-rknn/resources/coco_80_labels_list.txt";

    // 线程池初始化
    auto yolo_thread_pool = std::make_unique<YoloThreadPool>(model_path, 16,
                                                             labels_path);
    VideoFile video_file(input_path.c_str());
    Preprocess image_preprocess(video_file.get_frame_width(), video_file.get_frame_height(), 640);

    std::unique_ptr<cv::Mat> image;
    std::shared_ptr<cv::Mat> image_res;
    uint8_t running_flag = 0;
    cv::namedWindow("Video", cv::WINDOW_AUTOSIZE);
    static int image_count = 0;
    static int image_res_count = 0;

    int delay = 1000 / 25;

    TimeDuration time_duration;
    // yolo_thread_pool->AddInferenceTask(std::move(image), image_preprocess);
    // while (image_res == nullptr)
    // {
    //     image_res = yolo_thread_pool->GetImageResultFromQueue();
    // }
    // cv::imshow("Image demo", *image_res);
    // cv::waitKey(0);
    do
    {
        /* code */
        auto func = [&]
        {
            running_flag = 0;
            image = video_file.GetNextFrame();

            if (image != nullptr)
            {
                yolo_thread_pool->AddInferenceTask(std::move(image), image_preprocess);
                running_flag |= 0x01;
                image_count++;
            }

            image_res = yolo_thread_pool->GetImageResultFromQueue();

            if (image_res != nullptr)
            {
                cv::imshow("Video", *image_res);
                image_res_count++;
                NN_LOG_INFO("image count: %d, image res count: %d", image_count, image_res_count);
                cv::waitKey(1);
                running_flag |= 0x10;
            }
        };
        run_once_with_delay(func, std::chrono::milliseconds(delay));
    }
    // 读取图像或者获取结果非空 线程池里还有任务 推理的结果没有等于输入图像数
    while (running_flag || yolo_thread_pool->GetTaseksSize() ||
           image_count > image_res_count);

    auto time = std::chrono::duration_cast<std::chrono::milliseconds>(time_duration.DurationSinceLastTime());
    double fps = image_res_count * 1000.0 / time.count();
    NN_LOG_INFO("Total processed image count: %d", image_res_count);
    NN_LOG_INFO("Total time: %f ms, fps: %f", (double)time.count(), fps);

    yolo_thread_pool.reset();
    NN_LOG_INFO("Exit main");
    cv::destroyAllWindows();

    return 0;
}
